import matplotlib.pyplot as plt
from matplotlib.dates import DateFormatter
import numpy as np
import pandas as pd


class Plotter:
    @staticmethod
    def plot_analysis(data, signals, trend_lines):
        """绘制完整的分析结果图表"""
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(16, 10),
                                       gridspec_kw={'height_ratios': [3, 1]})

        # 价格图表
        ax1.plot(data.index, data['Close'],
                 label='Close Price', color='black', linewidth=1.5)

        # 绘制趋势线和信号
        Plotter._plot_trend_lines(ax1, trend_lines)
        Plotter._plot_signals(ax1, data, signals)

        # 成交量图表
        Plotter._plot_volume(ax2, data)

        # 设置图表格式
        ax1.set_title('Dow Theory Analysis', fontsize=16)
        ax1.legend(loc='upper left')
        ax1.grid(True, alpha=0.3)
        ax1.xaxis.set_major_formatter(DateFormatter("%Y-%m"))

        plt.tight_layout()
        return fig

    @staticmethod
    def _plot_trend_lines(ax, trend_lines):
        """绘制趋势线"""
        if trend_lines['primary_up']:
            line = trend_lines['primary_up']
            x = [line['start_date'], line['end_date']]
            y = [
                line['slope'] * x[0].toordinal() + line['intercept'],
                line['slope'] * x[1].toordinal() + line['intercept']
            ]
            ax.plot(x, y, 'g--', linewidth=1.5, label='Primary Up Trend')

        if trend_lines['primary_down']:
            line = trend_lines['primary_down']
            x = [line['start_date'], line['end_date']]
            y = [
                line['slope'] * x[0].toordinal() + line['intercept'],
                line['slope'] * x[1].toordinal() + line['intercept']
            ]
            ax.plot(x, y, 'r--', linewidth=1.5, label='Primary Down Trend')

    @staticmethod
    def _plot_signals(ax, data, signals):
        """绘制买卖信号"""
        if signals is None:
            return

        # 转换信号索引为日期类型
        signals.index = pd.to_datetime(signals.index)

        # 绘制买入信号（蓝色三角形）
        buy_signals = signals[signals['Signal'] == 1]
        if not buy_signals.empty:
            ax.scatter(buy_signals.index,
                       data.loc[buy_signals.index, 'Close'] * 0.98,
                       color='blue', s=100, marker='^',
                       label='Buy Signal', zorder=5)

        # 绘制卖出信号（红色三角形）
        sell_signals = signals[signals['Signal'] == -1]
        if not sell_signals.empty:
            ax.scatter(sell_signals.index,
                       data.loc[sell_signals.index, 'Close'] * 1.02,
                       color='red', s=100, marker='v',
                       label='Sell Signal', zorder=5)

    @staticmethod
    def _plot_volume(ax, data):
        """绘制成交量"""
        ax.bar(data.index, data['Volume'],
               color=np.where(data['Volume'] > data['Volume_MA'],
                              'green', 'gray'),
               alpha=0.6, width=0.6)
        ax.plot(data.index, data['Volume_MA'],
                color='red', linewidth=1, label='Volume MA(5)')
        ax.set_ylabel('Volume')
        ax.legend()
        ax.grid(True, alpha=0.3)